# 三维波动方程
# 我感觉用torch做矩阵运算比cupy还快一些...
# Gitee Repo

import torch
import numpy as np
import scipy
import time

torch.set_default_device('cuda' if torch.cuda.is_available() else 'cpu')

L  = 2;
dx = 0.05;
dt  = 0.001;
c=1

x,y,z=torch.meshgrid(torch.arange(-L,L,dx),torch.arange(-L,L,dx),torch.arange(-L,L,dx))
n  = x.shape[0]
sigma=(c*dt/dx)**2;

u0 = torch.zeros((n,n,n))
u1 = torch.zeros((n,n,n))
u2 = torch.zeros((n,n,n))

print(u0.device)
timer = time.time()

for tick in range(1000):
    print(tick)
    diff_i = u1[2:n,1:n-1,1:n-1]-2*u1[1:n-1,1:n-1,1:n-1]+u1[0:n-2,1:n-1,1:n-1]
    diff_j = u1[1:n-1,2:n,1:n-1]-2*u1[1:n-1,1:n-1,1:n-1]+u1[1:n-1,0:n-2,1:n-1]
    diff_k = u1[1:n-1,1:n-1,2:n]-2*u1[1:n-1,1:n-1,1:n-1]+u1[1:n-1,1:n-1,0:n-2]
    
    u2[1:n-1,1:n-1,1:n-1] = 2*u1[1:n-1,1:n-1,1:n-1] - u0[1:n-1,1:n-1,1:n-1]+sigma*(diff_i+diff_j+diff_k)
    u2[int(n/2),int(n/2),int(n/2)] += dt**2/dx**3*np.sin(2*np.pi/1*tick*dt);
    
    u0,u1,u2=u1,u2,u0
    
u_cpu = u1.cpu().numpy()[::2,::2,::2]
x_cpu = x.cpu().numpy() [::2,::2,::2]
y_cpu = y.cpu().numpy() [::2,::2,::2]
z_cpu = z.cpu().numpy() [::2,::2,::2]
print(time.time()-timer)

scipy.io.savemat('wave3D_data.mat',{'x':x_cpu,'y':y_cpu,'z':z_cpu,'u':u_cpu})
exit()
